import torch 
import torchvision
import random
import os

def num_compare(name:str):
    return int(name.split(".")[0])


def load_data_indexs(path="./g_seg_datas"):
    """
    从一个目录中加载数据集
    满足以下存储方式：数据为JPEG, PNG or GIF 格式，且image与label完全对应。
    ```
        ./path/
        ├── image
        └── label
    ```
    """
    datas = []
    image_path = os.path.join(path, "image")
    label_path = os.path.join(path, "label")
    f_names = os.listdir(image_path)
    f_names.sort(key=num_compare)
    for f_name in f_names:
        if f_name.endswith(".png"):
            if os.path.exists(os.path.join(image_path, f_name)):
                datas.append([os.path.join(image_path, f_name), os.path.join(label_path, f_name)])
    return datas

# 返回训练集和验证集
def get_dataset_spilt(path:str, rate:float=0.8, shuffle:bool=False):
    index_list = load_data_indexs(path)
    if shuffle:
        random.shuffle(index_list)
    split_point = int(len(index_list) * rate)
    return index_list[:split_point], index_list[split_point:]
  

class MyDataLoader():
    def __init__(self, input_data, batch_size:int=1, shuffle:bool=False, load_all=False, drop_last=False) -> None:
        super().__init__()
        if isinstance(input_data, str):
            self.data_path_index = load_data_indexs(input_data)
        elif isinstance(input_data, list):
            self.data_path_index = input_data
        else:
            raise Exception("input_data must be str or list")
        self.batch_size = batch_size
        self.load_all = load_all # TODO 是否加载所有数据进入硬盘
        if shuffle:
            random.shuffle(self.data_path_index)
        num_samples = len(self.data_path_index) // batch_size
        if len(self.data_path_index) % batch_size != 0:
            if not drop_last:
                num_samples += 1
        self.__batch_index = []
        count = 0
        for _ in range(num_samples):
            index_in_this_batch = []
            for _ in range(batch_size):
                if count < len(self.data_path_index):
                    index_in_this_batch.append(self.data_path_index[count])
                    count += 1
                else:
                    break
            self.__batch_index.append(index_in_this_batch.copy())

    def __getitem__(self, index):
        image = torchvision.io.read_image(self.data_path_index[index][0])
        label = torchvision.io.read_image(self.data_path_index[index][1])
        return image, label

    def __len__(self):
        return len(self.__batch_index)
    
    def __iter__(self):
        # 总循环
        for index_in_this_batch in self.__batch_index:
            image ,label = None, None
            # 每个batch循环
            for k, index in enumerate(index_in_this_batch):
                if image is None:
                    _image = torchvision.io.read_image(index[0])
                    _label = torchvision.io.read_image(index[1])
                    image = torch.zeros((self.batch_size, *_image.shape))
                    label = torch.zeros((self.batch_size, *_label.shape))
                    image[0,:,:] = _image
                    label[0,:,:] = _label
                else:
                    image[k, :, :] = torchvision.io.read_image(index[0])
                    label[k, :, :] = torchvision.io.read_image(index[1])
            yield image, label

    def get_batch_index(self):
        return self.__batch_index